#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
最终遗传算法实现
"""

import pandas as pd
import numpy as np
import random

def load_data():
    """加载数据"""
    df = pd.read_excel('附件1.xlsx')
    df.columns = ['sequence', 'config', 'color', 'drive']
    
    cars = []
    for _, row in df.iterrows():
        cars.append({
            'id': int(row['sequence']),
            'config': str(row['config']),
            'color': str(row['color']),
            'drive': str(row['drive'])
        })
    return cars

def calculate_fitness(sequence, cars_dict):
    """计算适应度"""
    # 颜色切换
    colors = [cars_dict[id]['color'] for id in sequence]
    switches = sum(1 for i in range(1, len(colors)) if colors[i] != colors[i-1])
    color_score = max(0, 100 - switches * 0.3)  # 降低惩罚
    
    # 驱动比例
    drives = [cars_dict[id]['drive'] for id in sequence]
    drive_deviation = 0
    for i in range(0, len(drives), 4):
        window = drives[i:i+4]
        if len(window) >= 2:
            two = window.count('两驱')
            four = window.count('四驱')
            ideal = len(window) / 2
            deviation = abs(two - ideal) + abs(four - ideal)
            drive_deviation += deviation
    drive_score = max(0, 100 - drive_deviation * 0.5)
    
    # 总得分
    total_score = color_score * 0.4 + drive_score * 0.3 + 85 * 0.2 + 95 * 0.1
    return total_score

def genetic_algorithm(cars, generations=20, population_size=20):
    """简化遗传算法"""
    cars_dict = {car['id']: car for car in cars}
    base_sequence = [car['id'] for car in cars]
    
    # 初始化种群
    population = []
    
    # 启发式初始解
    electric = [car['id'] for car in cars if car['color'] == '电动']
    combustion = [car['id'] for car in cars if car['color'] == '燃动']
    heuristic1 = electric + combustion
    
    two = [car['id'] for car in cars if car['drive'] == '两驱']
    four = [car['id'] for car in cars if car['drive'] == '四驱']
    heuristic2 = []
    i = j = 0
    while i < len(two) or j < len(four):
        if i < len(two):
            heuristic2.append(two[i])
            i += 1
        if j < len(four):
            heuristic2.append(four[j])
            j += 1
    
    # 添加初始解
    for seq in [base_sequence, heuristic1, heuristic2]:
        fitness = calculate_fitness(seq, cars_dict)
        population.append((seq, fitness))
    
    # 填充随机（确保数量足够）
    needed = max(0, population_size - len(population))
    for _ in range(needed):
        seq = base_sequence.copy()
        random.shuffle(seq)
        fitness = calculate_fitness(seq, cars_dict)
        population.append((seq, fitness))
    
    # 进化
    for gen in range(generations):
        population.sort(key=lambda x: x[1], reverse=True)
        
        if gen % 5 == 0:
            print(f"第{gen}代: 最佳={population[0][1]:.2f}")
        
        # 选择前60%
        parents = population[:max(2, population_size*3//5)]
        new_population = parents.copy()
        
        # 生成后代
        while len(new_population) < population_size:
            p1, p2 = random.sample(parents, 2)
            
            # 交叉
            child = p1[0].copy()
            
            # 简单的两点交叉
            if len(child) >= 3:
                i, j = sorted(random.sample(range(1, len(child)-1), 2))
                child[i:j] = p2[0][i:j]
            
            # 修复重复元素
            seen = set()
            duplicates = []
            for idx, val in enumerate(child):
                if val in seen:
                    duplicates.append(idx)
                else:
                    seen.add(val)
            
            # 填充缺失元素
            missing = [x for x in base_sequence if x not in seen]
            for idx, val in zip(duplicates, missing):
                child[idx] = val
            
            # 变异
            if random.random() < 0.1 and len(child) >= 2:
                i, j = random.sample(range(len(child)), 2)
                child[i], child[j] = child[j], child[i]
            
            fitness = calculate_fitness(child, cars_dict)
            new_population.append((child, fitness))
        
        population = new_population[:population_size]
    
    return population[0][0]

def create_schedule_matrix(sequence):
    """创建调度矩阵"""
    n_cols = max(100, (len(sequence) + 9) // 10)
    schedule = [[0] * n_cols for _ in range(10)]
    
    car_idx = 0
    for col in range(n_cols):
        for row in range(10):
            if car_idx < len(sequence):
                schedule[row][col] = sequence[car_idx]
                car_idx += 1
    return schedule

def save_results(cars, best_sequence, method):
    """保存结果"""
    # 调度矩阵
    schedule = create_schedule_matrix(best_sequence)
    pd.DataFrame(schedule).to_excel('result11.xlsx', index=False, header=False)
    
    # 车辆序列
    cars_dict = {car['id']: car for car in cars}
    sequence_data = [cars_dict[id] for id in best_sequence]
    pd.DataFrame(sequence_data).to_excel('result12.xlsx', index=False)
    
    print(f"使用{method}方法，结果已保存")

def main():
    cars = load_data()
    
    # 直接遗传算法
    best_sequence = genetic_algorithm(cars, generations=15, population_size=20)
    
    # 保存结果
    save_results(cars, best_sequence, "遗传算法")
    
    # 计算最终得分
    cars_dict = {car['id']: car for car in cars}
    colors = [cars_dict[id]['color'] for id in best_sequence]
    drives = [cars_dict[id]['drive'] for id in best_sequence]
    
    switches = sum(1 for i in range(1, len(colors)) if colors[i] != colors[i-1])
    print(f"最终颜色切换次数: {switches}")
    print(f"最终序列前10个: {best_sequence[:10]}")

if __name__ == "__main__":
    main()